#include <iostream>
#include <string>
#include <qpOASES.hpp>
#include "MPCsolver.h"

using namespace std;
using namespace Eigen; // conflict with namespace "qpOASES"
using qpOASES::QProblem;
using qpOASES::real_t;

MPCsolver::MPCsolver()
{
    qproblem = new qpOASES::QProblem(12 * PREDICTION_HORIZON, 20 * PREDICTION_HORIZON); // 优化变量维度nV, 约束变量维度nC

    Matrix<double, 13, 1> Q_tmpvec;
    Q_tmpvec << 25, 25, 10, 1, 1, 100, 0, 0, 0.3, 0.2, 0.2, 20, 0;
    Q_ = Q_tmpvec.replicate(PREDICTION_HORIZON, 1).asDiagonal(); // use diagonalMatrix to accelerate calculation

    Matrix<double, 12, 1> R_tmpvec;
    R_tmpvec << 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005;
    R_ = R_tmpvec.replicate(PREDICTION_HORIZON, 1).asDiagonal(); // use diagonalMatrix to accelerate calculation

    c_lb_.setZero(); // 0(20h*1)

    Matrix<double, 5, 1> c_i_ub;
    c_i_ub << INF, INF, INF, INF, f_max;
    c_ub_ = c_i_ub.replicate(PREDICTION_HORIZON, 1); // [INF, INF, INF, INF, f_max](20h*1)

    Matrix<double, 5, 3> C_i;
    C_i << -1, 0, miu,
        0, -1, miu,
        1, 0, miu,
        0, 1, miu,
        0, 0, 1;

    for (int i = 0; i < 4 * PREDICTION_HORIZON; i++)
    {
        C_.block<5, 3>(5 * i, 3 * i) = C_i;
    }
}

MPCsolver::~MPCsolver()
{
    delete qproblem;
}

void MPCsolver::setD_and_x0(const Matrix<double, 13 * PREDICTION_HORIZON, 1> &vecD, const Matrix<double, 13, 1> &x_0)
{
    D_ = vecD;
    x_0_ = x_0;
}

// void MPCsolver::calculateAll(const Matrix<double, PREDICTION_HORIZON, 1> &psi_k_d,
//                              const Matrix<double, 3 * PREDICTION_HORIZON, 4> &r_k_ix)
void MPCsolver::calculateAll(const Matrix<double, PREDICTION_HORIZON, 1> &psi_k_d,
                             const Matrix<double, 3 * PREDICTION_HORIZON, 4> &r_k_ix)
{
    // I_B数组映射为Eigen矩阵
    Matrix3d I_B = Map<Matrix3d>(I_B_array);

    int h = PREDICTION_HORIZON;
    double delt = DELTA_T_MPC;
    Matrix<double, 13, 13> A_k;
    Matrix3d tmpmat;
    // clang-format off
    A_k <<  1, 0, 0, 0, 0, 0, -1, -1, -1, 0, 0, 0, 0,
            0, 1, 0, 0, 0, 0, -1, -1, -1, 0, 0, 0, 0,
            0, 0, 1, 0, 0, 0, -1, -1, -1, 0, 0, 0, 0,
            0, 0, 0, 1, 0, 0, 0, 0, 0, 0, delt, 0, 0, 0,
            0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, delt, 0, 0,
            0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, delt, 0,
            0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, delt,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0;
    // clang-format on

    // 计算A(0)并赋值给A_qp
    A_k.block<3, 3>(0, 6) = (R("z", psi_k_d(0))).transpose() * delt;
    A_qp_.block<13, 13>(0, 0) = A_k;

    Matrix<double, 13, 12> B_k;
    B_k.setZero();

    Matrix3d R_k_B_O = R("z", psi_k_d(0));
    Matrix3d I_k_p_inverse = ((R_k_B_O * I_B * R_k_B_O.transpose()).inverse());

    //  计算B(0)并赋值给B_qp
    B_k.block<3, 3>(6, 0) = I_k_p_inverse * Antisymmetric(r_k_ix.col(0).segment(0, 3)) * delt;
    B_k.block<3, 3>(6, 3) = I_k_p_inverse * Antisymmetric(r_k_ix.col(1).segment(0, 3)) * delt;
    B_k.block<3, 3>(6, 6) = I_k_p_inverse * Antisymmetric(r_k_ix.col(2).segment(0, 3)) * delt;
    B_k.block<3, 3>(6, 9) = I_k_p_inverse * Antisymmetric(r_k_ix.col(3).segment(0, 3)) * delt;

    Matrix3d I3deltm = Matrix3d::Identity() * delt / Mass;
    for (int _i = 0; _i < 4; _i++)
        B_k.block<3, 3>(9, 3 * _i) = I3deltm;

    B_qp_.block<13, 12>(0, 0) = B_k;

    for (int i = 1; i < h; i++)
    {
        R_k_B_O = R("z", psi_k_d(i));
        I_k_p_inverse = ((R_k_B_O * I_B * R_k_B_O.transpose()).inverse());

        // 计算A(1)~A(h-1)并赋值给A_qp
        A_k.block<3, 3>(0, 6) = R_k_B_O.transpose() * delt;
        A_qp_.block<13, 13>(i * 13, 0) = A_k * A_qp_.block<13, 13>((i - 1) * 13, 0);

        // 计算B(1)~B(h-1)并赋值给B_qp
        B_k.block<3, 3>(6, 0) = I_k_p_inverse * Antisymmetric(r_k_ix.col(0).segment(0 + i * 3, 3)) * delt;
        B_k.block<3, 3>(6, 3) = I_k_p_inverse * Antisymmetric(r_k_ix.col(1).segment(0 + i * 3, 3)) * delt;
        B_k.block<3, 3>(6, 6) = I_k_p_inverse * Antisymmetric(r_k_ix.col(2).segment(0 + i * 3, 3)) * delt;
        B_k.block<3, 3>(6, 9) = I_k_p_inverse * Antisymmetric(r_k_ix.col(3).segment(0 + i * 3, 3)) * delt;

        Matrix3d I3deltm = Matrix3d::Identity() * delt / Mass;
        for (int _i = 0; _i < 4; _i++)
            B_k.block<3, 3>(9, 3 * _i) = I3deltm;

        B_qp_.block<13, 12>(13 * i, 12 * i) = B_k;

        for (int j = 0; j < i; j++)
        {
            B_qp_.block<13, 12>(13 * i, 12 * j) =
                A_k * B_qp_.block<13, 12>(13 * (i - 1), 12 * j);
        }
    }

    H_ = 2 * ((B_qp_.transpose() * Q_ * B_qp_) + R_.toDenseMatrix());
    g_ = 2 * B_qp_.transpose() * Q_ * (A_qp_ * x_0_ - D_);
}

void MPCsolver::qpSolveU_first()
{
    int h = PREDICTION_HORIZON;

    real_t *H_qp = H_.data();
    real_t *g_qp = g_.data(); // 线性项系数 与maatlab不同为 x' * g
    real_t *C_qp = C_.data(); // 不等式/等式约束系数矩阵
    real_t *lb_qp = NULL;
    real_t *ub_qp = NULL;
    real_t *lbC_qp = c_lb_.data(); // 不等式/等式上界 上下界相等时为等式约束
    real_t *ubC_qp = c_ub_.data(); // 不等式/等式下界

    real_t U_opt[12 * h];
    int nWSR = 50; // 最大迭代次数

    // cout << "before init" << endl;
    qproblem->init(H_qp, g_qp, C_qp, lb_qp, ub_qp, lbC_qp, ubC_qp, nWSR);
    qproblem->getPrimalSolution(U_opt);

    // cout << "----------------------f_MPC-------------------------" << endl;
    // cout << "f1: " << U_opt[0] << "   " << U_opt[1] << "   " << U_opt[2] << endl;
    // cout << "f2: " << U_opt[3] << "   " << U_opt[4] << "   " << U_opt[5] << endl;
    // cout << "f3: " << U_opt[6] << "   " << U_opt[7] << "   " << U_opt[8] << endl;
    // cout << "f4: " << U_opt[9] << "   " << U_opt[10] << "   " << U_opt[11] << endl;
}

void MPCsolver::qpSolveU_next(Matrix<double, 12, 1> &f_MPC)
{
    int h = PREDICTION_HORIZON;

    real_t *g_qp = g_.data(); // 线性项系数 与matlab不同为 x' * g

    real_t *lb_qp = NULL;
    real_t *ub_qp = NULL;
    real_t *lbC_qp = c_lb_.data(); // 不等式/等式上界 上下界相等时为等式约束
    real_t *ubC_qp = c_ub_.data(); // 不等式/等式下界

    real_t U_opt[12 * h];
    int nWSR = 50; // 最大迭代次数
    qproblem->hotstart(g_qp, lb_qp, ub_qp, lbC_qp, ubC_qp, nWSR);
    qproblem->getPrimalSolution(U_opt);

    f_MPC << U_opt[0], U_opt[1], U_opt[2], U_opt[3], U_opt[4], U_opt[5], U_opt[6], U_opt[7], U_opt[8], U_opt[9], U_opt[10], U_opt[11];

    // cout << "----------------------f_MPC-------------------------" << endl;
    // cout << "f1: " << U_opt[0] << "   " << U_opt[1] << "   " << U_opt[2] << endl;
    // cout << "f2: " << U_opt[3] << "   " << U_opt[4] << "   " << U_opt[5] << endl;
    // cout << "f3: " << U_opt[6] << "   " << U_opt[7] << "   " << U_opt[8] << endl;
    // cout << "f4: " << U_opt[9] << "   " << U_opt[10] << "   " << U_opt[11] << endl;
    // cout << "hotstart" << endl;
}
